﻿using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Diagnostics;
using System.Management.Automation;
using System.Management.Automation.Runspaces;
using System.Text;

namespace Wolfpack.Contrib.Checks.Powershell
{
    public class RunspaceController : IWolfpackPowerShellHost
    {
        private bool _disposed;
        private Runspace _runspace;
        private Pipeline _scriptExecutionPipeline;
        private bool _scriptRequestedStop;

        /// <summary>
        /// </summary>
        /// <remarks>
        /// The script that gets executed is wrapped in a function.
        /// The is named based on the Runspace.InstanceId
        /// </remarks>
        private readonly object _scriptLock = new object();

        /// <summary>
        /// Synchronizes write access to the console buffer.
        /// </summary>
        private readonly object _outputLock = new object();

        /// <summary>
        /// Stores contents that are written to the console.
        /// </summary>
        private readonly StringBuilder _consoleOutput = new StringBuilder();

        /// <summary>
        /// 
        /// </summary>
        private bool _stopRequested;

        /// <summary>
        /// Gets the console output.
        /// </summary>
        public string ConsoleOutput
        {
            get
            {
                lock (_outputLock)
                {
                    string output = _consoleOutput.ToString();
                    return output;
                }
            }
        }

        /// <summary>
        /// Initializes the <see cref="Runspace"/> with a <see cref="WolfpackPSHost"/>.
        /// </summary>
        /// <param name="factory">The configuration factory.</param>
        public void Initialize(IRunspaceConfigurationFactory factory)
        {
            try
            {
                WolfpackPSHost host = new WolfpackPSHost(this);
                RunspaceConfiguration runspaceConfiguration = factory.Create();
                _runspace = RunspaceFactory.CreateRunspace(host, runspaceConfiguration);
            }
            catch (Exception)
            {
                // TODO: determine the kind of exceptions that can be thrown
                throw;
            }
        }

        /// <summary>
        /// Runs the script.
        /// </summary>
        /// <param name="scriptName">Name of the script.</param>
        /// <param name="scriptBody">The script body.</param>
        /// <param name="parameters">The parameters.</param>
        /// <returns>The PipelineResults</returns>
        public PipelineResults RunScript(string scriptName, string scriptBody, Dictionary<string, string> parameters)
        {
            if (_runspace == null)
            {
                throw new InvalidOperationException("The Initialize method must be called before executing a script.");
            }

            Pipeline functionDeclarationPipeline = null;
            Pipeline defaultOutputPipeline = null;

            try
            {
                lock (_scriptLock)
                {
                    if (_stopRequested)
                    {
                        return PipelineResults.Empty;
                    }

                    // create pipeline to create the script body as a function
                    string functionName = "__" + _runspace.InstanceId.ToString("N");
                    functionDeclarationPipeline = CreateFunctionDeclarationPipeline(scriptBody, functionName);

                    // create pipeline to execute function
                    _scriptExecutionPipeline = CreateScriptExecutionPipeline(parameters, functionName);

                    // create pipeline to send the output to the default formatter and to the default output cmdlet.
                    defaultOutputPipeline = _runspace.CreatePipeline();
                    defaultOutputPipeline.Commands.Add("out-default");

                    _runspace.Open();

                    if (_stopRequested)
                    {
                        return PipelineResults.Empty;
                    }

                    // since the script supplied is user-input, it could cause
                    // a parse exception
                    try
                    {
                        functionDeclarationPipeline.Invoke();
                    }
                    catch (ParseException exception)
                    {
                        return new PipelineResults(exception);
                    }
                }

                // there is a chance that the caller has already requested the pipeline be stopped
                Collection<PSObject> pSObjects = _scriptExecutionPipeline.Invoke();

                defaultOutputPipeline.Input.Write(pSObjects);
                defaultOutputPipeline.Invoke();

                if (_scriptExecutionPipeline.PipelineStateInfo.State == PipelineState.Completed || _scriptRequestedStop)
                {
                    lock (_outputLock)
                    {
                        // wait for all of the output to be written
                    }
                }

                if (_scriptExecutionPipeline.PipelineStateInfo.State == PipelineState.Stopped && _stopRequested)
                {
                    return new PipelineResults(_scriptExecutionPipeline.PipelineStateInfo.Reason);
                }

                // get the results from the pipeline
                return new PipelineResults(_scriptExecutionPipeline, pSObjects);
            }
            finally
            {
                DisposePipeline(ref defaultOutputPipeline);
                DisposePipeline(ref _scriptExecutionPipeline);
                DisposePipeline(ref functionDeclarationPipeline);
            }
        }

        private void DisposePipeline(ref Pipeline pipeline)
        {
            if (pipeline == null)
            {
                return;
            }

            pipeline.Dispose();
            pipeline = null;
        }

        /// <summary>
        /// Stops the execution of the script.
        /// </summary>
        public void StopScript()
        {
            lock (_scriptLock)
            {
                _stopRequested = true;
                // copy to local variable to avoid race condition
                Pipeline pipeline = _scriptExecutionPipeline;
                if (pipeline != null)
                {
                    pipeline.Stop();
                }
            }
        }

        /// <summary>
        /// Request to end the current runspace. The Windows PowerShell 
        /// engine calls this method to request that the host application 
        /// shut down and terminate the host root runspace. 
        /// </summary>
        /// <param name="exitCode">
        /// The exit code that is used to set the host's process exit code.
        /// </param>
        public void SetShouldExit(int exitCode)
        {
            _scriptRequestedStop = true;
            // copy to local variable to avoid race condition
            Pipeline pipeline = _scriptExecutionPipeline;
            if (pipeline != null)
            {
                pipeline.StopAsync();
            }
        }

        /// <summary>
        /// Writes characters to the output display of the host.
        /// </summary>
        /// <param name="value">The characters to be written.</param>
        public void Write(string value)
        {
            lock (_outputLock)
            {
                _consoleOutput.Append(value);
            }
        }

        /// <summary>
        /// Writes characters to the output display of the host with possible foreground and background colors. 
        /// </summary>
        /// <param name="foregroundColor">The foreground color of the display.</param>
        /// <param name="backgroundColor">The background color of the display.</param>
        /// <param name="value">The characters to be written.</param>
        public void Write(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value)
        {
            lock (_outputLock)
            {
                _consoleOutput.Append(value);
            }
        }

        /// <summary>
        /// Displays a debug message to the user. 
        /// </summary>
        /// <param name="message">Debug message to be displayed.</param>
        public void WriteDebugLine(string message)
        {
            Trace.WriteLine(message);
        }

        public void WriteErrorLine(string value)
        {
            lock (_outputLock)
            {
                _consoleOutput.AppendLine(value);
            }
        }

        /// <summary>
        /// Writes a new line character (carriage return) to the output display of the host.
        /// </summary>
        public void WriteLine()
        {
            lock (_outputLock)
            {
                _consoleOutput.AppendLine();
            }
        }

        /// <summary>
        /// Writes a new line character (carriage return) to the output display of the host.
        /// </summary>
        /// <param name="value">The line of characters to be written.</param>
        public void WriteLine(string value)
        {
            lock (_outputLock)
            {
                _consoleOutput.AppendLine(value);
            }
        }

        /// <summary>
        /// Writes characters to the output display of the host with possible foreground and background colors. 
        /// </summary>
        /// <param name="foregroundColor">The foreground color of the display.</param>
        /// <param name="backgroundColor">The background color of the display.</param>
        /// <param name="value">The characters to be written.</param>
        public void WriteLine(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value)
        {
            lock (_outputLock)
            {
                _consoleOutput.AppendLine(value);
            }
        }

        /// <summary>
        /// Writes a progress report to be displayed to the user. 
        /// </summary>
        /// <param name="sourceId">
        /// Unique identifier of the source of the record. An int64 is used because the 
        /// 'this' pointer of the command where the record originated from is used, and 
        /// that may be from a remote runspace on a 64-bit computer.
        /// </param>
        /// <param name="record">
        /// A <see cref="ProgressRecord"/> object that contains the progress record to be displayed.
        /// </param>
        /// <remarks>
        /// This method is invoked by a cmdlet call to the <see cref="WriteProgress"/> method.
        /// </remarks>
        public void WriteProgress(long sourceId, ProgressRecord record)
        {
            if (record != null)
            {
                string message = string.Format("Progress Record - Source Id = {0}, Content = {1}", sourceId, record);
                Trace.WriteLine(message);
            }
        }

        /// <summary>
        /// Writes a verbose line to be displayed to the user. 
        /// </summary>
        /// <param name="message">The verbose message to be displayed.</param>
        public void WriteVerboseLine(string message)
        {
            Trace.WriteLine(message);
        }

        /// <summary>
        /// Writes a warning line to be displayed to the user. 
        /// </summary>
        /// <param name="message">The warning message to be displayed.</param>
        public void WriteWarningLine(string message)
        {
            lock (_outputLock)
            {
                _consoleOutput.AppendLine(message);
            }

            if (!string.IsNullOrEmpty(message))
            {
                Trace.WriteLine(message);
            }
        }

        /// <summary>
        /// Releases unmanaged and - optionally - managed resources
        /// </summary>
        public void Dispose()
        {
            // copy to local variable to avoid race condition
            Pipeline pipeline = _scriptExecutionPipeline;

            if (pipeline != null)
            {
                pipeline.Stop();
            }

            Dispose(true);
            GC.SuppressFinalize(this);
        }

        private Pipeline CreateScriptExecutionPipeline(Dictionary<string, string> parameters, string functionName)
        {
            Pipeline pipeline = _runspace.CreatePipeline();
            Command scriptExecutionCommand = new Command(functionName, false);

            // add the parameters to the script excution command
            foreach (var parameter in parameters)
            {
                scriptExecutionCommand.Parameters.Add(parameter.Key, parameter.Value);
            }

            // merge the errors stream to the output stream
            scriptExecutionCommand.MergeMyResults(PipelineResultTypes.Error, PipelineResultTypes.Output);
            pipeline.Commands.Add(scriptExecutionCommand);

            return pipeline;
        }

        private Pipeline CreateFunctionDeclarationPipeline(string scriptBody, string functionName)
        {
            Pipeline functionDeclarationPipeline = _runspace.CreatePipeline();

            string functionDeclaration = string.Format("function {0} {{ {1} }}", functionName, scriptBody);
            Command functionDeclarationCommand = new Command(functionDeclaration, true);

            functionDeclarationPipeline.Commands.Add(functionDeclarationCommand);
            return functionDeclarationPipeline;
        }

        /// <summary>
        /// Releases unmanaged and - optionally - managed resources
        /// </summary>
        /// <param name="disposing">
        /// <c>true</c> to release both managed and unmanaged resources; 
        /// <c>false</c> to release only unmanaged resources.
        /// </param>
        private void Dispose(bool disposing)
        {
            if (!_disposed && disposing)
            {
                if (_scriptExecutionPipeline != null)
                {
                    _scriptExecutionPipeline.Dispose();
                    _scriptExecutionPipeline = null;
                }

                if (_runspace != null)
                {
                    _runspace.Dispose();
                    _runspace = null;
                }
            }

            _disposed = true;
        }
    }
}